import json
import math
import os
from PIL import Image

import numpy as np


class LineBenchmark():
    def __init__(self, generate_dataset_images_path, gt_label_path, pred_label_path=None):
        self.outlier = None
        self.rec = None
        self.total = None
        self.total_score = None
        self.rec_line = None
        self.total_line = None
        self.scores = None
        self.images_path = generate_dataset_images_path
        self.gt_path = gt_label_path
        self.pred_path = pred_label_path
        self.gt = dict()
        self.pred = dict()
        self.load()

    def load(self):
        with open(self.gt_path, "r") as fp:
            gt = json.load(fp)

        for key in gt.keys():
            labels = gt[key]
            if "Line" in labels.keys():
                self.gt[key] = labels["Line"]

        with open(self.pred_path, "r") as fp:
            pred_data = json.load(fp)

        for absolute_path in pred_data:
            v = pred_data[absolute_path]
            filename = absolute_path.split("\\")[-1]
            self.pred[filename] = v

    @staticmethod
    def getMasksYByX(masks, x, y, h):
        """
        masks 是直线的掩码图，包含多条直线
        x 是采样x坐标
        """
        def getMaskYIndexByX(mask, x):
            # 二分查找
            l_index = 0
            r_index = len(mask) - 1

            xrange = mask[-1][0] - mask[0][0]

            if mask[l_index][0] > x:
                    return 0

            if mask[r_index][0] < x:
                    return -1

            m_index = (l_index + r_index) / 2
            m_index = int(m_index)
            while l_index < r_index:
                if l_index + 1 == r_index and mask[l_index][0] < x < mask[r_index][0]:
                    break

                elif mask[m_index][0] == x:
                    break

                elif mask[m_index][0] < x:
                    l_index = m_index
                    m_index = (l_index + r_index) / 2
                    m_index = int(m_index)
                    continue

                else:
                    r_index = m_index
                    m_index = (l_index + r_index) / 2
                    m_index = int(m_index)
                    continue
            return m_index

        ys = []
        diff = []
        for mask in masks:
            index = getMaskYIndexByX(mask, x)
            if index is None:
                ys.append(0)
            else:
                ys.append(mask[index][1])
            diff.append(abs(ys[-1] - y) / h)

        return ys, diff

    @staticmethod
    def check(xs, ys, h, masks, threshold):
        # 得到所有的masks对应的ys
        ys_e = []
        diff_e = []
        for x_index in range(len(xs)):
            _ys, _diff = LineBenchmark.getMasksYByX(masks, xs[x_index], ys[x_index], h)
            ys_e.append(_ys)
            diff_e.append(_diff)

        diff_e = np.array(diff_e)
        ys_e = np.array(ys_e)

        # 求差距的均值
        mean_diff_e = np.mean(diff_e, axis=0)
        if len(mean_diff_e.shape) == 0 or mean_diff_e.shape[0] == 0:
            return False, None

        min_index = np.argmin(mean_diff_e)

        # 求
        if mean_diff_e[min_index] > threshold:
            return False, None

        return True, np.array(ys_e)[:, min_index]

    def run(self, threshold=0.05):
        total = 0
        rec = 0
        outlier = 0
        total_score = 0
        # 取每个gt与pred做对比
        for key in self.gt.keys():
            image_path = os.path.join(self.images_path, key)
            img = Image.open(image_path)
            w, h = img.size
            img.close()

            gt = self.gt[key]
            masks = self.pred[key]

            for item in gt:
                # 取gt的采样点对所有mask进行采样，确定最小方差的mask匹配
                _gt_xs = item[::2]
                _gt_ys = item[1::2]
                gt_ys = []
                gt_xs = []
                # gt_ys不能有等于0的情况
                for index in range(len(_gt_ys)):
                    tmp_gt_y = _gt_ys[index]
                    if tmp_gt_y <= 1:
                        outlier += 1
                        continue
                    else:
                        gt_ys.append(tmp_gt_y)
                        gt_xs.append(_gt_xs[index])

                # 预测点总数
                total += len(gt_xs)
                # 采样过滤
                isPass, pred_ys = self.check(gt_xs, gt_ys, h, masks, threshold=threshold)

                if isPass:
                    rec += len(gt_xs)
                    gt_ys = np.array(gt_ys)
                    pred_ys = np.array(pred_ys)
                    diff = np.abs(pred_ys - gt_ys) / h
                    diff = np.ones_like(diff) - diff
                    score = np.sum(diff)
                    total_score += score

        self.total = total
        self.rec = rec
        self.total_score = total_score
        self.outlier = outlier

    def summary(self):
        print(self.outlier, self.rec, self.total, self.total_score)
        print(f"识别率：{round(self.rec/self.total * 100, 4)}")
        print(f"召回率：{round(self.total_score/self.total * 100, 4)}")

if __name__ == "__main__":
    gt_labels = r"D:\Code\script\_banchmark\benchmark\dataset\CV\line\labels\val.json"
    pred_labels = r"D:\Code\script\_banchmark\benchmark\utils\lineformer_pred.json"
    benchmark = LineBenchmark(gt_labels, pred_labels)
    benchmark.run()
    benchmark.summary()
    # [[theta_y, theta, [center_x, center_y], R, string_center, color, [x1, y1, x2, y2]]]
